#include "fastrt/reid_interface.h"

int main()
{
    REIDParams params;
    params.weights_path = "../../sbs_R50-ibn.wts";
    params.engine_path = "../../sbs_R50-ibn.engine.trt";
    params.device_id = 1;
    params.max_batch_size = 10;
    params.inputW = 128;
    params.inputH = 384;

    FastREID reid;
    reid.build(&params);

    std::vector <cv::Mat> images;
    for (size_t b = 0; b < 21; b++)
    {
        images.emplace_back(cv::Mat(params.inputW, params.inputH, CV_8UC3, cv::Scalar(255, 255, 255)));
    }

    std::vector<std::vector<float> > feat_embedding;

    /* run inference */
    TimePoint start_infer, end_infer;
    int LOOP_TIMES = 100;
    start_infer = Time::now();
    for (int times = 0; times < LOOP_TIMES; ++times)
    {
        feat_embedding.clear();
        int flag = reid.infer(images, feat_embedding);
    }
    end_infer = Time::now();

    /* print output */
    std::cout << "feat_embedding: " << feat_embedding.size() << std::endl;
    std::cout << "feat_embedding: " << feat_embedding[0].size() << std::endl;
    for (size_t img_idx = 0; img_idx < feat_embedding.size(); ++img_idx) {
        std::cout << "img_idx: " << img_idx << " ";
        for (int dim = 0; dim < feat_embedding[0].size(); ++dim) {
            std::cout << feat_embedding[img_idx][dim] << " ";
            if ((dim + 1) % 10 == 0) {
                std::cout << std::endl;
            }
        }
        std::cout << std::endl;
    }
    std::cout << std::endl;

    std::cout << "[Preprocessing+Inference]: " <<
              std::chrono::duration_cast<std::chrono::milliseconds>(end_infer - start_infer).count() /
              static_cast<float>(LOOP_TIMES) << "ms" << std::endl;

    return 0;
}

